from ._core import NdArray, Dtype, DtypeTag
from tinynumpy import _core
import numbers
import numpy as np
import ctypes

# 定义类型别名
float32 = Dtype(DtypeTag.FP32)
float64 = Dtype(DtypeTag.FP64)
int32 = Dtype(DtypeTag.INT32)
bool_ = Dtype(DtypeTag.BOOL)


def convert_value(val, dtype):
    if dtype == float32:
        return ctypes.c_float(val)
    elif dtype == float64:
        return ctypes.c_double(val)  # Python的float对应C++的double
    elif dtype == int32:
        return ctypes.c_int(val)
    elif dtype == bool_:
        # 支持多种布尔值表示
        if isinstance(val, str):
            val = val.lower()
            if val in ("true", "t", "yes", "y", "1"):
                return True
            elif val in ("false", "f", "no", "n", "0"):
                return False
        return ctypes.c_bool(val)
    else:
        raise ValueError(f"Unsupported dtype: {dtype}")


class ndarray:
    def __init__(self, obj=None, dtype=None, shape=None, copy=True):
        """
        创建新的 ndarray

        参数:
            obj: 输入数据 (列表、元组、NumPy 数组等)
            dtype: 数据类型 (float32, float64, int32, bool_)
            shape: 指定形状
            copy: 是否复制数据
        """
        if dtype is None:
            dtype = float32

        if obj is None:
            if shape is None:
                self._arr = NdArray(dtype, [])
            else:
                self._arr = NdArray(dtype, shape)
        elif isinstance(obj, NdArray):
            self._arr = obj.copy() if copy else obj
        elif isinstance(obj, np.ndarray):
            self._arr = self._from_numpy(obj, dtype)
        elif isinstance(obj, (list, tuple)):
            self._arr = self._from_nested(obj, dtype)
        elif isinstance(obj, numbers.Number):
            self._arr = NdArray(dtype, [1])
            self._arr[0] = float(obj)
        else:
            raise TypeError(f"Unsupported input type: {type(obj)}")

    def _compute_flat_index(index, shape):
        """计算多维索引对应的扁平化索引"""
        if len(index) != len(shape):
            raise ValueError(
                f"Index length {len(index)} does not match shape length {len(shape)}"
            )

        flat_index = 0
        stride = 1
        for i in range(len(shape) - 1, -1, -1):
            dim = index[i]
            if dim < 0:
                dim = shape[i] + dim
            if dim < 0 or dim >= shape[i]:
                raise IndexError(
                    f"Index {dim} out of range for dimension {i} with size {shape[i]}"
                )
            flat_index += dim * stride
            stride *= shape[i]

        return flat_index

    @staticmethod
    def _from_nested(obj, dtype):
        """从嵌套列表创建数组"""

        def flatten(nested_list):
            def _get_shape(nested):
                """递归获取嵌套列表的形状"""
                if not isinstance(nested, (list, tuple)):
                    return ()  # 标量，形状为空元组

                if len(nested) == 0:
                    return (0,)  # 空列表，形状为(0,)

                # 检查所有子元素的形状是否一致
                shapes = set()
                for item in nested:
                    shapes.add(_get_shape(item))

                if len(shapes) > 1:
                    raise ValueError("All sublists must have the same shape")

                # 当前维度的长度 + 子元素的形状
                sub_shape = shapes.pop() if shapes else ()
                return (len(nested),) + sub_shape

            # 辅助函数：递归扁平化
            def _flatten(lst):
                result = []
                for item in lst:
                    if isinstance(item, list):
                        result.extend(_flatten(item))
                    else:
                        result.append(item)
                return result

            # 主逻辑
            try:
                shape = _get_shape(nested_list)  # 尝试获取形状
            except (TypeError, ValueError):
                shape = None  # 不规则嵌套返回None

            return _flatten(nested_list), shape

        flat_data, shape = flatten(obj)

        # 创建底层数组
        arr = NdArray(dtype, shape)
        for i, val in enumerate(flat_data):
            arr[i] = convert_value(val, dtype)

        return arr

    @staticmethod
    def _from_numpy(np_arr, dtype=None):
        """从 NumPy 数组创建"""
        if dtype is None:
            if np_arr.dtype == np.float32:
                dtype = float32
            elif np_arr.dtype == np.float64:
                dtype = float64
            elif np_arr.dtype == np.int32:
                dtype = int32
            elif np_arr.dtype == np.bool_:
                dtype = bool_
            else:
                dtype = float32  # 默认

        # 创建底层数组
        arr = NdArray(dtype, np_arr.shape)
        flat = np_arr.flatten()
        for i in range(len(flat)):
            arr[i] = flat[i]
        return arr

    @property
    def shape(self):
        return self._arr.shape

    @property
    def dtype(self):
        return self._arr.dtype

    @property
    def ndim(self):
        return self._arr.ndim

    @property
    def size(self):
        return self._arr.size

    def __getitem__(self, index):
        if isinstance(index, int):
            return self._arr[index]
        elif isinstance(index, tuple):
            return self._arr[index]
        elif index is Ellipsis:
            return self
        else:
            raise TypeError(f"Unsupported index type: {type(index)}")

    def __setitem__(self, index, value):
        if isinstance(index, int):
            self._arr[index] = value
        elif isinstance(index, tuple):
            self._arr[index] = value
        else:
            raise TypeError(f"Unsupported index type: {type(index)}")

    def __repr__(self):
        return f"ndarray({self._arr.to_string()}, dtype={self.dtype})"

    def __str__(self):
        return self.__repr__()

    def tolist(self):
        """转换为嵌套列表"""

        def build_nested(shape, data, offset=0):
            if len(shape) == 1:
                return list(data[offset : offset + shape[0]])

            result = []
            for i in range(shape[0]):
                size = np.prod(shape[1:])
                result.append(build_nested(shape[1:], data, offset + i * size))
            return result

        data = [self._arr[i] for i in range(self.size)]
        return build_nested(self.shape, data)

    def copy(self):
        return ndarray(self._arr.copy())

    def reshape(self, new_shape):
        """改变数组形状"""
        arr = self._arr.reshape(new_shape)
        return ndarray(arr)

    def astype(self, dtype):
        """转换数据类型"""
        arr = self._arr.astype(dtype)
        return ndarray(arr)

    def sum(self, axis=None, keepdims=False):
        """沿指定轴求和"""
        return ndarray(self._arr.sum(axis, keepdims))

    # 添加更多 NumPy 风格的 API
    def mean(self, axis=None, keepdims=False):
        """沿指定轴求平均值"""
        if axis is None:
            total = self.sum()
            return total / self.size
        else:
            raise NotImplementedError("Axis-specific mean not implemented yet")

    def __add__(self, other):
        if isinstance(other, ndarray):
            return ndarray(self._arr + other._arr)
        elif isinstance(other, numbers.Number):
            result = self.copy()
            for i in range(result.size):
                result[i] += other
            return result
        else:
            raise TypeError(f"Unsupported type for addition: {type(other)}")

    def __sub__(self, other):
        if isinstance(other, ndarray):
            return ndarray(self._arr - other._arr)
        elif isinstance(other, numbers.Number):
            result = self.copy()
            for i in range(result.size):
                result[i] -= other
            return result
        else:
            raise TypeError(f"Unsupported type for subtraction: {type(other)}")

    def __mul__(self, other):
        if isinstance(other, ndarray):
            return ndarray(self._arr * other._arr)
        elif isinstance(other, numbers.Number):
            result = self.copy()
            for i in range(result.size):
                result[i] *= other
            return result
        else:
            raise TypeError(f"Unsupported type for multiplication: {type(other)}")

    def __truediv__(self, other):
        if isinstance(other, ndarray):
            return ndarray(self._arr / other._arr)
        elif isinstance(other, numbers.Number):
            result = self.copy()
            for i in range(result.size):
                result[i] /= other
            return result
        else:
            raise TypeError(f"Unsupported type for division: {type(other)}")

    def __neg__(self):
        result = self.copy()
        for i in range(result.size):
            result[i] = -result[i]
        return result

    def __eq__(self, other):
        if isinstance(other, ndarray):
            return self._arr == other._arr
        elif isinstance(other, numbers.Number):
            result = ndarray(shape=self.shape, dtype=bool_)
            for i in range(self.size):
                result[i] = self[i] == other
            return result
        else:
            raise TypeError(f"Unsupported type for comparison: {type(other)}")

    def to_numpy(self, copy=False):
        """
        将ndarray对象转换为NumPy数组

        参数:
            copy: 是否复制数据（默认False，共享内存）

        返回:
            numpy.ndarray: 包含相同数据的NumPy数组
        """
        # 处理空数组
        if self.size == 0:
            return np.empty(self.shape, dtype=self._get_numpy_dtype())

        # 获取底层数据指针
        data_ptr = self._arr.data_ptr

        # 确定NumPy数据类型
        np_dtype = self._get_numpy_dtype()

        # 计算总字节数
        nbytes = self.size * self._arr.itemsize

        # 创建共享内存的NumPy数组
        ctypes_buffer = (ctypes.c_byte * nbytes).from_address(data_ptr)
        np_arr = np.frombuffer(ctypes_buffer, dtype=np_dtype).reshape(self.shape)

        # 设置正确的步长（以元素为单位）
        strides = tuple(s // self._arr.itemsize for s in self._arr.strides)
        if strides != np_arr.strides:
            np_arr = np.lib.stride_tricks.as_strided(
                np_arr, shape=self.shape, strides=strides
            )

        # 如果需要复制数据，创建一个副本
        if copy:
            return np_arr.copy()

        # 设置数组为可写
        # np_arr.flags.writeable = True

        # 确保原始数组的生命周期与NumPy数组绑定
        # np_arr.base = self

        return np_arr

    def _get_numpy_dtype(self):
        """获取对应的NumPy数据类型"""
        if self.dtype.tag == DtypeTag.FP32:
            return np.float32
        elif self.dtype.tag == DtypeTag.FP64:
            return np.float64
        elif self.dtype.tag == DtypeTag.INT32:
            return np.int32
        elif self.dtype.tag == DtypeTag.BOOL:
            return np.bool_
        else:
            raise TypeError(f"Unsupported dtype: {self.dtype}")


# 工厂函数
def array(obj, dtype=None):
    return ndarray(obj, dtype)


def zeros(shape, dtype=float32):
    arr = NdArray.zeros(dtype, shape)
    return ndarray(arr)


def ones(shape, dtype=float32):
    arr = NdArray.ones(dtype, shape)
    return ndarray(arr)


def empty(shape, dtype=float32):
    arr = NdArray.empty(dtype, shape)
    return ndarray(arr)


def uniform(low, high, shape, dtype=float32):
    arr = _core.uniform(convert_value(low, dtype).value, convert_value(high, dtype).value, shape)
    return ndarray(arr)
